# -*- coding: utf-8 -*-
# backend/utils/rate_limiter.py
import time
from typing import Dict, Optional
from collections import defaultdict, deque


class RateLimiter:
    """速率限制器"""
    
    def __init__(self):
        # 存储每个key的请求时间戳
        self.requests: Dict[str, deque] = defaultdict(deque)
    
    async def check_rate_limit(
        self, 
        key: str, 
        max_requests: int, 
        window_seconds: int
    ) -> bool:
        """检查是否超过速率限制"""
        
        current_time = time.time()
        window_start = current_time - window_seconds
        
        # 获取该key的请求队列
        request_queue = self.requests[key]
        
        # 移除窗口外的旧请求
        while request_queue and request_queue[0] < window_start:
            request_queue.popleft()
        
        # 检查当前请求数是否超过限制
        if len(request_queue) >= max_requests:
            return False
        
        # 添加当前请求时间戳
        request_queue.append(current_time)
        return True
    
    def reset_key(self, key: str):
        """重置特定key的速率限制"""
        if key in self.requests:
            del self.requests[key]
    
    def get_remaining_requests(self, key: str, max_requests: int, window_seconds: int) -> int:
        """获取剩余可用请求数"""
        current_time = time.time()
        window_start = current_time - window_seconds
        
        request_queue = self.requests[key]
        
        # 清理过期请求
        while request_queue and request_queue[0] < window_start:
            request_queue.popleft()
        
        return max(0, max_requests - len(request_queue))
    
    def get_reset_time(self, key: str, window_seconds: int) -> Optional[float]:
        """获取速率限制重置时间"""
        request_queue = self.requests.get(key)
        if not request_queue:
            return None
        
        return request_queue[0] + window_seconds